'''
Author: error: error: git config user.name & please set dead value or install git && error: git config user.email & please set dead value or install git & please set dead value or install git
Date: 2025-09-13 23:59:32
LastEditors: error: error: git config user.name & please set dead value or install git && error: git config user.email & please set dead value or install git & please set dead value or install git
LastEditTime: 2025-09-29 23:14:17
FilePath: /ml-pro/project/hmm_model.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''
import numpy as np
import pandas as pd
from pathlib import Path
from hmmlearn import hmm
from tqdm import tqdm

# ---------- 4. 加载「红+蓝」联合观测序列 ----------
def load_red_blue_obs_sequences(data_dir: str):
    """
    返回 List[np.ndarray]  每条长度=7，值域 1~34
    前 6 个是排序后的红球，第 7 个是篮球（用 34 表示）
    """
    files = Path(data_dir).glob('*.csv')
    obs_seqs = []
    for f in files:
        df = pd.read_csv(f)
        # 红球
        df['red_sorted'] = (
            df['红球'].str.split()
                    .apply(lambda x: sorted(map(int, x)))
        )
        # 篮球
        df['blue_int'] = df['蓝球'].astype(int)
        # 组装成 7 个符号
        for _, row in df.iterrows():
            seq = list(row['red_sorted']) + [34]   # 34 代表篮球
            obs_seqs.append(np.array(seq, dtype=int))
    return obs_seqs


# ---------- 5. 训练「红+蓝」联合 HMM ----------
def fit_redblue_hmm(obs_seqs, n_hidden=8, n_iter=50, random_state=42):
    """
    输入：List[np.ndarray] 每条长度=7
    返回：训练好的 CategoricalHMM
    """
    # 构造适合CategoricalHMM的格式
    X = np.concatenate(obs_seqs).reshape(-1, 1)
    lengths = [7] * len(obs_seqs)  # 每个序列长度为7

    model = hmm.CategoricalHMM(
                n_components=n_hidden,
                n_iter=n_iter,
                random_state=random_state,
                init_params='ste',          # 初始化起始概率、转移概率和发射概率
                params='ste',               # 训练过程中更新起始概率、转移概率和发射概率
                verbose=True)
    # 添加全面的平滑处理，避免概率为0
    model.transmat_prior = 1.0  # 转移矩阵的先验值
    model.startprob_prior = 1.0  # 起始概率的先验值
    model.emissionprob_prior = 1.0  # 发射概率的先验值
    model.fit(X, lengths)
    return model


# ---------- 6. 预测下一期「红+蓝」 ----------
def predict_next_redblue(model):
    # 使用CategoricalHMM的sample方法预测下一期
    sample, _ = model.sample(7)  # 预测7个符号（6红+1蓝）
    # 提取红球和蓝球
    red_balls = sorted([int(x) for x in sample[:6, 0]])
    blue_ball = int(sample[6, 0])
    
    # 确保红球在1-33范围内，蓝球在1-16范围内
    red_balls = [min(max(1, ball), 33) for ball in red_balls]
    blue_ball = min(max(1, blue_ball), 16)
    
    return red_balls, blue_ball


# ---------- 7. 仅预测篮球（可选） ----------
def load_blue_only_sequences(data_dir: str):
    """只返回篮球序列，用于单独建模"""
    files = Path(data_dir).glob('*.csv')
    blue_seqs = []
    for f in files:
        df = pd.read_csv(f)
        blue_seqs += df['蓝球'].astype(int).tolist()
    return [np.array([b], dtype=int) for b in blue_seqs]


def fit_blue_only_hmm(blue_seqs, n_hidden=6, n_iter=50, random_state=42):
    """
    篮球值域 1-16，用 MultinomialHMM，n_trials=1
    """
    X = []
    for seq in blue_seqs:
        cnt = np.zeros(16, dtype=int)
        cnt[seq[0] - 1] = 1
        X.append(cnt)
    X = np.array(X)
    lengths = [1] * len(blue_seqs)

    model = hmm.CategoricalHMM(
                n_components=n_hidden,
                n_iter=n_iter,
                random_state=random_state,
                init_params='ste',          # 初始化起始概率、转移概率和发射概率
                params='ste',               # 训练过程中更新起始概率、转移概率和发射概率
                verbose=True)
    # 添加全面的平滑处理，避免概率为0
    model.transmat_prior = 1.0  # 转移矩阵的先验值
    model.startprob_prior = 1.0  # 起始概率的先验值
    model.emissionprob_prior = 1.0  # 发射概率的先验值
    model.fit(X, lengths)
    return model


def predict_next_blue(model):
    sample, _ = model.sample(1)
    # CategoricalHMM返回的是观测符号索引（从0开始）
    symbol_idx = sample[0, 0]
    return symbol_idx + 1      # 1-16


# ---------- 8. 主流程（红+蓝联合版） ----------
if __name__ == '__main__':
    data_dir = 'data'

    # 方案 A：红+蓝一起预测
    obs_seqs_rb = load_red_blue_obs_sequences(data_dir)
    print(f'共 {len(obs_seqs_rb)} 期数据（红+蓝）')
    # 减少隐藏状态数量，避免过拟合
    model_rb = fit_redblue_hmm(obs_seqs_rb, n_hidden=3, n_iter=50)
    for i in range(5):
        red, blue = predict_next_redblue(model_rb)
        print(f'联合预测 {i+1}: 红球 {red}  篮球 {blue}')

    # # 方案 B：篮球单独预测（可选）
    # blue_seqs = load_blue_only_sequences(data_dir)
    # model_b = fit_blue_only_hmm(blue_seqs, n_hidden=6, n_iter=100)
    # for i in range(5):
    #     print(f'单独篮球预测 {i+1}: 篮球 {predict_next_blue(model_b)}')